import os
import random
import string
import base64

from gmssl import sm4, sm2


class SM4(object):
    def __init__(self):
        self.crypt_sm4 = sm4.CryptSM4()
        self.crypt_sm4_iv = base64.b64decode("AQ4Zvt54xKn9QaW86ZzWdg==") # 这里请填写你的iv值

    # sm4 cbc 加密
    def encrypt(self, encrypt_key, value):
        """
        国密sm4加密
        :param encrypt_key: sm4加密key
        :param value: 待加密的字符串
        :return: sm4加密后的十六进制值
        """
        crypt_sm4 = self.crypt_sm4
        crypt_sm4.set_key(encrypt_key.encode(), sm4.SM4_ENCRYPT)
        date_str = str(value)
        encrypt_value = crypt_sm4.crypt_cbc(self.crypt_sm4_iv, date_str.encode())
        # 返回base64编码的加密数据
        return base64.b64encode(encrypt_value).decode()

    # sm4 cbc 解密
    def decrypt(self, decrypt_key, encrypt_value):
        """
        国密sm4解密
        :param decrypt_key:sm4加密key
        :param encrypt_value: 待解密的十六进制值
        :return: 原字符串
        """
        encrypted_bytes = base64.b64decode(encrypt_value)
        crypt_sm4 = self.crypt_sm4
        crypt_sm4.set_key(decrypt_key.encode(), sm4.SM4_DECRYPT)  # 设置密钥
        decrypt_value = crypt_sm4.crypt_cbc(self.crypt_sm4_iv, encrypted_bytes)  # 开始解密。十六进制类型
        return decrypt_value.decode()

    # 生成128位随机密钥key
    @staticmethod
    def generate_key():
        """
        生成128位随机密钥key
        :return: 128位随机密钥key
        """
        return ''.join(random.sample(string.ascii_letters + string.digits, 16))


class SM2(object):
    def __init__(self, helipay_public_key, public_key = None, private_key = None):
        self.public_key = public_key
        self.private_key = private_key
        self.helipay_public_key = helipay_public_key

    # 公钥加密
    def helipay_encrypt(self, data, java=False):
        sm2_crypt = sm2.CryptSM2(public_key=self.helipay_public_key, private_key='', mode=1, asn1=True)
        encrypt_res = sm2_crypt.encrypt(data.encode("UTF_8"))
        if java:
            encrypt_res = b'\x04' + encrypt_res
        return base64.b64encode(encrypt_res).decode("utf-8")

    def encrypt(self, data, java=False):
        """sm2加密方法"""
        sm2_crypt = sm2.CryptSM2(public_key=self.public_key, private_key='', mode=1, asn1=True)
        encrypt_res = sm2_crypt.encrypt(data.encode("UTF_8"))
        if java:
            encrypt_res = b'\x04' + encrypt_res
        return base64.b64encode(encrypt_res).decode("utf-8")

    def decrypt(self, ciphertext: str, java=True):
        """sm2解密方法"""
        encrypted_bytes = base64.b64decode(ciphertext)
        if java:
            encrypted_bytes = encrypted_bytes[1:]  # 移除0x04前缀
        
        sm2_crypt = sm2.CryptSM2(public_key='', private_key=self.private_key, mode=1, asn1=True)
        decrypted_bytes = sm2_crypt.decrypt(encrypted_bytes)
        return decrypted_bytes.decode("utf-8")

    # 私钥签名
    def generate_sign(self, data):
        sm2_crypt = sm2.CryptSM2(private_key=self.private_key, public_key=self.public_key, mode=1, asn1=True)
        random_hex_str = os.urandom(32).hex()  # 随机数
        signature = sm2_crypt.sign_with_sm3(data.encode("utf-8"), random_hex_str)
        return base64.b64encode(bytes.fromhex(signature)).decode('utf-8')

    # 公钥验签
    def verify_sign(self, sign, data):
        sm2_crypt = sm2.CryptSM2(private_key=None, public_key=self.public_key, mode=1, asn1=True)
        return sm2_crypt.verify(sign, data)


